package api

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"gitee.com/linxing_3/sr-core/domain/api"
	"gitee.com/linxing_3/sr-core/plugins/log"
	"gitee.com/mkwchecking/sr-admin/apps/knowledge/acl"
	"gitee.com/mkwchecking/sr-admin/apps/knowledge/service"
	"github.com/gin-gonic/gin"
	llmApi "github.com/ollama/ollama/api"
)

func NewRagApi(ragSvc *service.RAGService) *RAGApi {
	return &RAGApi{
		ragSvc: ragSvc,
	}
}

type RAGApi struct {
	ragSvc *service.RAGService
}

func (r *RAGApi) Chat(c *gin.Context) {
	req := &acl.RAGChatReq{}

	apiCtx := api.NewCtx(c).WithReq(req)

	if !apiCtx.Do() {
		return
	}

	streamMod := req.Stream != nil && *req.Stream
	msgChan := make(chan service.ChatChan)
	// defer close(msgChan)

	var err error
	_, _, err = r.ragSvc.Chat(apiCtx.Context(), req, msgChan)

	if err != nil {
		apiCtx.WriteErr(nil, err, err.Error())
		return
	} else if streamMod {
		r.chatStreamResponse(apiCtx, c, msgChan)
		return
	} else {
		var result *llmApi.ChatResponse
		for msg := range msgChan {
			if msg.Err != nil {
				apiCtx.WriteErr(nil, err, msg.Err.Error())
				return
			} else if result == nil {
				nmsg := msg.Msg
				result = &nmsg
			} else {
				result.Message.Content += msg.Msg.Message.Content
				fmt.Print(msg.Msg.Message.Content)
			}
		}

		apiCtx.WriteCustom(result, http.StatusOK)
	}
}

func (r *RAGApi) chatStreamResponse(ctx api.IApiCtx, c *gin.Context, ch chan service.ChatChan) {
	c.Header("Content-Type", "application/x-ndjson")
	c.Stream(func(w io.Writer) bool {
		val, ok := <-ch
		if !ok {
			return false
		} else if val.Err != nil {
			return false
		}
		// fmt.Printf(val.Msg.Message.Content)
		bts, _ := json.Marshal(val.Msg)
		if _, err := w.Write(append(bts, '\n')); err != nil {
			log.Errorf("write stream error: %s", err)
			return false
		}
		return true
	})
}
